#include "web_server.h"
#include <atomic>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <iostream>
#include <sstream>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <signal.h>
#include <unistd.h>
#include <cerrno>
#include <cstring>

// Ensure full payload is sent over TCP
static bool send_all(int fd, const char* data, size_t len) {
    size_t sent = 0;
    while (sent < len) {
        ssize_t n = ::send(fd, data + sent, len - sent, MSG_NOSIGNAL);
        if (n < 0) {
            if (errno == EINTR) continue;
            int e = errno;
            std::cerr << "send_all: send failed after " << sent << "/" << len
                      << " bytes, errno=" << e << " (" << strerror(e) << ")" << std::endl;
            return false;
        }
        if (n == 0) {
            std::cerr << "send_all: peer closed while sending after " << sent
                      << "/" << len << " bytes" << std::endl;
            return false;
        }
        sent += static_cast<size_t>(n);
    }
    return true;
}

struct WebMJPEGServer::Impl {
    int port;
    std::atomic<bool> running{false};

    std::thread server_thread;
    std::function<std::vector<uint8_t>()> frame_provider;
    std::function<std::string()> stats_provider;
    std::function<std::string()> time_provider;
    std::function<std::string()> stats_json_provider;
    std::atomic<int> client_count{0};
};

WebMJPEGServer::WebMJPEGServer(int port) : impl_(new Impl{port}) {}
WebMJPEGServer::~WebMJPEGServer() { stop(); delete impl_; }

void WebMJPEGServer::setFrameProvider(std::function<std::vector<uint8_t>()> provider) {
    impl_->frame_provider = std::move(provider);
}

void WebMJPEGServer::setStatsProvider(std::function<std::string()> stats_html_provider,
                                      std::function<std::string()> time_provider) {
    impl_->stats_provider = std::move(stats_html_provider);
    impl_->time_provider = std::move(time_provider);
}

void WebMJPEGServer::setStatsJsonProvider(std::function<std::string()> stats_json_provider) {
    impl_->stats_json_provider = std::move(stats_json_provider);
}

void WebMJPEGServer::start() {
    if (impl_->running) return;
    impl_->running = true;
    // Prevent process termination on write to closed socket
    signal(SIGPIPE, SIG_IGN);
    impl_->server_thread = std::thread([this](){
        int server_fd = socket(AF_INET, SOCK_STREAM, 0);
        int opt = 1;
        setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
        sockaddr_in addr{};
        addr.sin_family = AF_INET;
        addr.sin_addr.s_addr = INADDR_ANY;
        addr.sin_port = htons(impl_->port);
        bind(server_fd, (struct sockaddr*)&addr, sizeof(addr));
        listen(server_fd, 4);
        while (impl_->running) {
            sockaddr_in cliaddr{}; socklen_t clilen = sizeof(cliaddr);
            int client_fd = accept(server_fd, (struct sockaddr*)&cliaddr, &clilen);
            if (client_fd < 0) continue;
            char ipbuf[INET_ADDRSTRLEN] = {0};
            inet_ntop(AF_INET, &cliaddr.sin_addr, ipbuf, sizeof(ipbuf));
            uint16_t cport = ntohs(cliaddr.sin_port);
            std::cout << "clientid = " << client_fd << ", Accepted from " << ipbuf << ":" << cport << std::endl;
            // set small recv timeout to avoid blocking on dead connections
            timeval rcv_to{.tv_sec=2, .tv_usec=0};
            setsockopt(client_fd, SOL_SOCKET, SO_RCVTIMEO, &rcv_to, sizeof(rcv_to));
            timeval snd_to{.tv_sec=2, .tv_usec=0};
            setsockopt(client_fd, SOL_SOCKET, SO_SNDTIMEO, &snd_to, sizeof(snd_to));
            char req_buf[1024] = {0};
            size_t total = 0;
            bool read_ok = false;
            while (total < sizeof(req_buf)-1) {
                ssize_t rn = recv(client_fd, req_buf + total, sizeof(req_buf)-1 - total, 0);
                if (rn > 0) {
                    total += static_cast<size_t>(rn);
                    req_buf[total] = '\0';
                    // stop once headers complete
                    if (std::string(req_buf, total).find("\r\n\r\n") != std::string::npos) {
                        read_ok = true;
                        break;
                    }
                    // continue reading until timeout or buffer full
                    continue;
                } else if (rn == 0) {
                    std::cout << "Request read: peer closed immediately, client " << client_fd
                              << " from " << ipbuf << ":" << cport << std::endl;
                    break;
                } else { // rn < 0
                    int e = errno;
                    if (e == EAGAIN || e == EWOULDBLOCK) {
                        std::cout << "Request read timeout (errno=" << e << ": " << strerror(e)
                                  << ") from " << ipbuf << ":" << cport << ", closing client " << client_fd << std::endl;
                    } else {
                        std::cout << "Request read failed (errno=" << e << ": " << strerror(e)
                                  << ") from " << ipbuf << ":" << cport << ", closing client " << client_fd << std::endl;
                    }
                    break;
                }
            }
            if (!read_ok && total == 0) { close(client_fd); continue; }
            std::string req_str(req_buf, total);
            std::cout << "Request: " << req_str.substr(0, req_str.find('\n')) << std::endl;
            if (req_str.find("GET /stats_time") == 0 || req_str.find("GET /stats_time ") == 0) {
                std::string html = impl_->time_provider ? impl_->time_provider() : std::string("-");
                std::ostringstream oss;
                oss << "HTTP/1.0 200 OK\r\nConnection: close\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: " << html.size() << "\r\n\r\n";
                {
                    auto s = oss.str();
                    send_all(client_fd, s.c_str(), s.size());
                }
                send_all(client_fd, html.c_str(), html.size());
                close(client_fd);
            } else if (req_str.find("GET /stats_json") == 0 || req_str.find("GET /stats_json ") == 0) {
                std::string json = impl_->stats_json_provider ? impl_->stats_json_provider() : std::string("{\"fps\":0,\"objects\":0,\"people\":0,\"vehicles\":0}");
                std::cout << "Serve /stats_json, bytes=" << json.size() << std::endl;
                std::ostringstream oss;
                oss << "HTTP/1.0 200 OK\r\nConnection: close\r\nContent-Type: application/json; charset=utf-8\r\nContent-Length: " << json.size() << "\r\n\r\n";
                {
                    auto s = oss.str();
                    send_all(client_fd, s.c_str(), s.size());
                }
                send_all(client_fd, json.c_str(), json.size());
                close(client_fd);
            } else if (req_str.find("GET /stats") == 0 || req_str.find("GET /stats ") == 0) {
                std::string html = impl_->stats_provider ? impl_->stats_provider() : std::string("<div/>\n");
                std::ostringstream oss;
                oss << "HTTP/1.0 200 OK\r\nConnection: close\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: " << html.size() << "\r\n\r\n";
                {
                    auto s = oss.str();
                    send_all(client_fd, s.c_str(), s.size());
                }
                send_all(client_fd, html.c_str(), html.size());
                close(client_fd);
            } else if (req_str.find("GET /mjpeg") == 0 || req_str.find("GET /mjpeg ") == 0) {
                std::cout << "/mjpeg connect start" << std::endl;
                std::thread([this, client_fd]() {
                    impl_->client_count.fetch_add(1, std::memory_order_relaxed);
                    std::cout << "client_count++ => " << impl_->client_count.load() << std::endl;
                    std::string header =
                        "HTTP/1.0 200 OK\r\n"
                        "Server: MJPEGStreamer/1.0\r\n"
                        "Content-Type: multipart/x-mixed-replace; boundary=frame\r\n"
                        "Cache-Control: no-cache\r\n"
                        "Pragma: no-cache\r\n\r\n";
                    send_all(client_fd, header.c_str(), header.size());
                    while (impl_->running) {
                        if (!impl_->frame_provider) { usleep(10000); continue; }
                        auto jpeg = impl_->frame_provider();
                        if (jpeg.empty()) { usleep(10000); continue; }
                        std::ostringstream oss;
                        oss << "--frame\r\nContent-Type: image/jpeg\r\nContent-Length: " << jpeg.size() << "\r\n\r\n";
                        {
                            auto s = oss.str();
                            if (!send_all(client_fd, s.c_str(), s.size())) break;
                        }
                        if (!send_all(client_fd, reinterpret_cast<const char*>(jpeg.data()), jpeg.size())) break;
                        if (!send_all(client_fd, "\r\n", 2)) break;
                        usleep(40000);
                    }
                    close(client_fd);
                    impl_->client_count.fetch_sub(1, std::memory_order_relaxed);
                    std::cout << "client_count-- => " << impl_->client_count.load() << std::endl;
                }).detach();
            } else if (req_str.find("GET /test") == 0 || req_str.find("GET /test ") == 0) {
                std::cout << "Serve /test" << std::endl;
                std::string body = "OK\n";
                std::ostringstream oss;
                oss << "HTTP/1.0 200 OK\r\nConnection: close\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: "
                    << body.size() << "\r\n\r\n";
                auto s = oss.str();
                send_all(client_fd, s.c_str(), s.size());
                send_all(client_fd, body.c_str(), body.size());
                close(client_fd);
            } else if (req_str.find("GET /favicon.ico") == 0) {
                const char* hdr = "HTTP/1.0 204 No Content\r\nConnection: close\r\nContent-Length: 0\r\n\r\n";
                send_all(client_fd, hdr, strlen(hdr));
                close(client_fd);
            } else if (req_str.find("GET / ") == 0 || req_str.find("GET /HTTP") == 0) {
                std::cout << "Serve index /" << std::endl;
                std::string html = std::string(R"(
<!DOCTYPE html>
<html>
<head>
  <title>YOLO Web</title>
  <meta charset='utf-8'>
  <meta name='viewport' content='width=device-width,initial-scale=1'>
  <style>
    body { background:#181c24; color:#f3f6fa; margin:0; font-family:'Segoe UI',Arial,sans-serif; }
    .title-bar { width:50vw; min-width:320px; margin:36px auto 18px auto; font-size:2.0em; font-weight:700; text-align:center; letter-spacing:2px; background: linear-gradient(90deg,#2d3a4a 0%,#1e90ff 100%); color:#fff; border-radius:18px; box-shadow:0 4px 24px #0008; padding:18px 0; user-select:none; }
    .stats-row { display:flex; justify-content:center; gap:18px; margin:0 auto; max-width:70vw; min-width:320px; }
    .stat-card { background:#223044; border-radius:12px; padding:14px 18px; text-align:center; box-shadow:0 2px 12px #0006; min-width:120px; }
    .stat-label { font-size:14px; color:#aabbd0; }
    .stat-value { font-size:24px; font-weight:700; color:#8cb3f4; margin-top:6px; }
    .video-area { display:flex; flex-direction:column; align-items:center; margin-top:18px; }
    .video-frame { max-width:70vw; min-width:320px; border-radius:14px; box-shadow:0 0 16px #000b; background:#111; margin-bottom:12px; }
    .video-label { color:#aaa; font-size:1em; margin-bottom:24px; }
  </style>
</head>
<body>
  <div class='title-bar'>YOLO 实时检测</div>
  <div class='stats-row'>
    <div class='stat-card'><div class='stat-label'>FPS</div><div class='stat-value' id='fps'>0.0</div></div>
    <div class='stat-card'><div class='stat-label'>检测目标数</div><div class='stat-value' id='obj_count'>0</div></div>
    <div class='stat-card'><div class='stat-label'>人员</div><div class='stat-value' id='person_count'>0</div></div>
    <div class='stat-card'><div class='stat-label'>车辆</div><div class='stat-value' id='car_count'>0</div></div>
  </div>
  <div class='video-area'>
    <img class='video-frame' src='/mjpeg'>
    <div class='video-label'>实时检测视频</div>
  </div>
  <div id='update_time' style='text-align:center; color:#8cb3f4; font-size:14px; margin:10px 0 30px 0;'>更新时间：<span id='time_text'>-</span></div>
  <script>
    function updateStats() {
      fetch('/stats_json').then(r => r.json()).then(data => {
        document.getElementById('fps').textContent = data.fps.toFixed(1);
        document.getElementById('obj_count').textContent = data.objects;
        document.getElementById('person_count').textContent = data.people;
        document.getElementById('car_count').textContent = data.vehicles;
      }).catch(()=>{});
      fetch('/stats_time').then(r2 => r2.text()).then(t => {
        document.getElementById('time_text').textContent = t;
      }).catch(()=>{});
    }
    setInterval(updateStats, 1000);
    updateStats();
  </script>
</body>
</html>
)");
                std::ostringstream oss;
                oss << "HTTP/1.0 200 OK\r\nConnection: close\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: " << html.size() << "\r\n\r\n";
                {
                    auto s = oss.str();
                    send_all(client_fd, s.c_str(), s.size());
                }
                send_all(client_fd, html.c_str(), html.size());
                close(client_fd);
            } else {
                // default 404
                std::string body = "<h3>404 Not Found</h3>";
                std::ostringstream oss;
                oss << "HTTP/1.0 404 Not Found\r\nConnection: close\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: " << body.size() << "\r\n\r\n";
                auto s = oss.str();
                send_all(client_fd, s.c_str(), s.size());
                send_all(client_fd, body.c_str(), body.size());
                close(client_fd);
            }
        }
        close(server_fd);
    });
}

void WebMJPEGServer::stop() {
    if (!impl_->running) return;
    impl_->running = false;
    if (impl_->server_thread.joinable()) impl_->server_thread.join();
}

bool WebMJPEGServer::isRunning() const { return impl_->running; }

int WebMJPEGServer::getClientCount() const { return impl_->client_count.load(std::memory_order_relaxed); }
